import sys
import os
sys.path.append("/home/cole/project_graduate/mmdetection")
from mmdet.apis.inference import init_detector,inference_detector
from mmengine.config import Config
import torch
if __name__ == '__main__':
    config_file = "/home/cole/project_graduate/mmdetection/configs/yolox/yolox_tiny_8xb8-300e_coco.py"
    cfg = Config.fromfile(config_file)
    model_path = "/home/cole/project_graduate/mmdetection/model/yolox_tiny_8x8_300e_coco_20211124_171234-b4047906.pth"
    model = init_detector(cfg,model_path)
    model.eval()
    input_names = ["image"]
    output_names = ["S_cls","M_cls","L_cls","S_bbox","M_bbox","L_bbox","S_conf","M_conf","L_conf"]
    input_tensor = torch.zeros(1,3,416,416,device=next(model.parameters()).device)
    torch.onnx.export(model,input_tensor,"/home/cole/project_graduate/mmdetection/model/yolox.onnx",input_names=input_names, output_names=output_names)
